/**********************************************************************
 *  This program is free software; you can redistribute it and/or     *
 *  modify it under the terms of the GNU General Public License       *
 *  as published by the Free Software Foundation; either version 2    *
 *  of the License, or (at your option) any later version.            *
 *                                                                    *
 *  This program is distributed in the hope that it will be useful,   *
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of    *
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the     *
 *  GNU General Public License for more details.                      *
 *                                                                    *
 *  You should have received a copy of the GNU General Public License *
 *  along with this program; if not, see http://gnu.org/licenses/     *
 *  ---                                                               *
 *  Copyright (C) 2009, Justin Davis <tuxdavis@gmail.com>             *
 *  Copyright (C) 2009-2017 ImageWriter developers                    *
 *                 https://sourceforge.net/projects/win32diskimager/  *
 **********************************************************************/

#ifndef WINVER
#define WINVER 0x0601
#endif

#include "pch.h"
#include "disk.h"

CString GetErrorMessage(LPCTSTR format)
{
    UINT nLastError = GetLastError();
    LPTSTR szError = NULL;
    ::FormatMessage(FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_ALLOCATE_BUFFER, NULL, nLastError, 0,
        (LPTSTR)&szError, 0, NULL);

    CString szFormat(format);
    szFormat.Append(TEXT("\nError %d: %s"));

    CString szMsg;
    szMsg.Format((LPCTSTR)szFormat, nLastError, szError);
    LocalFree(szError);
    return szMsg;
}

HANDLE GetHandleOnFile(LPCWSTR filelocation, DWORD access)
{
    HANDLE hFile;
    hFile = CreateFileW(filelocation, access, (access == GENERIC_READ) ? FILE_SHARE_READ : 0, NULL, (access == GENERIC_READ) ? OPEN_EXISTING:CREATE_ALWAYS, 0, NULL);
    if(hFile == INVALID_HANDLE_VALUE)
    {
        CString szMsg = GetErrorMessage(TEXT("An error occurred when attempting to get a handle on the file."));
        ::MessageBox(NULL, (LPCTSTR)szMsg, TEXT("File Error"), MB_OK);
    }
    return hFile;
}
DWORD GetDeviceID(HANDLE hVolume)
{
    VOLUME_DISK_EXTENTS sd;
    DWORD dwBytesReturned;
    if(!DeviceIoControl(hVolume, IOCTL_VOLUME_GET_VOLUME_DISK_EXTENTS, NULL, 0, &sd, sizeof(sd), &dwBytesReturned, NULL))
    {
        CString szMsg = GetErrorMessage(TEXT("An error occurred when attempting to get information on volume."));
        ::MessageBox(NULL, (LPCTSTR)szMsg, TEXT("Volume Error"), MB_OK);
    }
    return sd.Extents[0].DiskNumber;
}

HANDLE GetHandleOnDevice(int device, DWORD access)
{
    HANDLE hDevice;
    CString szDeviceName;
    szDeviceName.Format(_T("\\\\.\\PhysicalDrive%d"), device);
    hDevice = CreateFile(szDeviceName, access, FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, OPEN_EXISTING, 0, NULL);
    if(hDevice == INVALID_HANDLE_VALUE)
    {
        CString szMsg = GetErrorMessage(TEXT("An error occurred when attempting to get a handle on the device."));
        ::MessageBox(NULL, (LPCTSTR)szMsg, TEXT("Device Error"), MB_OK);
    }
    return hDevice;
}

HANDLE GetHandleOnVolume(int volume, DWORD access)
{
    HANDLE hVolume;
    char szVolumeName[] = "\\\\.\\A:";
    szVolumeName[4] += volume;
    hVolume = CreateFileA(szVolumeName, access, FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, OPEN_EXISTING, 0, NULL);
    if(hVolume == INVALID_HANDLE_VALUE)
    {
        CString szMsg = GetErrorMessage(TEXT("An error occurred when attempting to get a handle on the volume."));
        ::MessageBox(NULL, (LPCTSTR)szMsg, TEXT("Volume Error"), MB_OK);
    }
    return hVolume;
}

BOOL GetLockOnVolume(HANDLE handle)
{
    DWORD bytesreturned;
    BOOL bResult;
    bResult = DeviceIoControl(handle, FSCTL_LOCK_VOLUME, NULL, 0, NULL, 0, &bytesreturned, NULL);
    if(!bResult)
    {
        CString szMsg = GetErrorMessage(TEXT("An error occurred when attempting to lock the volume."));
        ::MessageBox(NULL, (LPCTSTR)szMsg, TEXT("Lock Error"), MB_OK);
    }
    return (bResult);
}

BOOL RemoveLockOnVolume(HANDLE handle)
{
    DWORD junk;
    BOOL bResult;
    bResult = DeviceIoControl(handle, FSCTL_UNLOCK_VOLUME, NULL, 0, NULL, 0, &junk, NULL);
    if(!bResult)
    {
        CString szMsg = GetErrorMessage(TEXT("An error occurred when attempting to unlock the volume."));
        ::MessageBox(NULL, (LPCTSTR)szMsg, TEXT("Unlock Error"), MB_OK);
    }
    return (bResult);
}

BOOL UnmountVolume(HANDLE handle)
{
    DWORD junk;
    BOOL bResult;
    bResult = DeviceIoControl(handle, FSCTL_DISMOUNT_VOLUME, NULL, 0, NULL, 0, &junk, NULL);
    if(!bResult)
    {
        CString szMsg = GetErrorMessage(TEXT("An error occurred when attempting to dismount the volume."));
        ::MessageBox(NULL, (LPCTSTR)szMsg, TEXT("Dismount Error"), MB_OK);
    }
    return (bResult);
}

BOOL IsVolumeUnmounted(HANDLE handle)
{
    DWORD junk;
    BOOL bResult;
    bResult = DeviceIoControl(handle, FSCTL_IS_VOLUME_MOUNTED, NULL, 0, NULL, 0, &junk, NULL);
    return (!bResult);
}

char* ReadSectorDataFromHandle(HANDLE handle, unsigned long long startSector, unsigned long long numSectors, unsigned long long sectorSize)
{
    unsigned long bytesRead;
    char *data = new char[sectorSize * numSectors];
    LARGE_INTEGER li;
    li.QuadPart = startSector * sectorSize;
    SetFilePointer(handle, li.LowPart, &li.HighPart, FILE_BEGIN);
    if(!ReadFile(handle, data, sectorSize * numSectors, &bytesRead, NULL))
    {
        CString szMsg = GetErrorMessage(TEXT("An error occurred when attempting to read data from handle."));
        ::MessageBox(NULL, (LPCTSTR)szMsg, TEXT("Read Error"), MB_OK);

        delete[] data;
        data = NULL;
    }
    if(data && bytesRead < (sectorSize * numSectors))
    {
        memset(data + bytesRead,0,(sectorSize * numSectors) - bytesRead);
    }
    return data;
}

BOOL WriteSectorDataToHandle(HANDLE handle, char* data, unsigned long long startSector, unsigned long long numSectors, unsigned long long sectorSize)
{
    unsigned long bytesWritten;
    BOOL bResult;
    LARGE_INTEGER li;
    li.QuadPart = startSector * sectorSize;
    SetFilePointer(handle, li.LowPart, &li.HighPart, FILE_BEGIN);
    bResult = WriteFile(handle, data, sectorSize * numSectors, &bytesWritten, NULL);
    if(!bResult)
    {
        CString szMsg = GetErrorMessage(TEXT("An error occurred when attempting to write data to handle."));
        ::MessageBox(NULL, (LPCTSTR)szMsg, TEXT("Write Error"), MB_OK);
    }
    return bResult;
}

unsigned long long GetNumberOfSectors(HANDLE handle, unsigned long long *sectorSize)
{
    DWORD junk;
    DISK_GEOMETRY_EX diskgeometry;
    BOOL bResult;
    bResult = DeviceIoControl(handle, IOCTL_DISK_GET_DRIVE_GEOMETRY_EX, NULL, 0, &diskgeometry, sizeof(diskgeometry), &junk, NULL);
    if(!bResult)
    {
        CString szMsg = GetErrorMessage(TEXT("An error occurred when attempting to get the device's geometry."));
        ::MessageBox(NULL, (LPCTSTR)szMsg, TEXT("Device Error"), MB_OK);
        return 0;
    }
    if(sectorSize != NULL)
    {
        *sectorSize = (unsigned long long)diskgeometry.Geometry.BytesPerSector;
    }
    return (unsigned long long)diskgeometry.DiskSize.QuadPart / (unsigned long long)diskgeometry.Geometry.BytesPerSector;
}

unsigned long long GetFileSizeInSectors(HANDLE handle, unsigned long long sectorSize)
{
    unsigned long long retVal = 0;
    if(sectorSize) // avoid divide by 0
    {
        LARGE_INTEGER filesize;
        if(GetFileSizeEx(handle, &filesize) == 0)
        {
            retVal = 0;
            CString szMsg = GetErrorMessage(TEXT("An error occurred while getting the file size."));
            ::MessageBox(NULL, (LPCTSTR)szMsg, TEXT("File Error"), MB_OK);
        }
        else
        {
            retVal = ((unsigned long long)filesize.QuadPart / sectorSize ) + (((unsigned long long)filesize.QuadPart % sectorSize) ? 1 : 0);
        }
    }
    return(retVal);
}

BOOL SpaceAvailable(char* location, unsigned long long spaceneeded)
{
    ULARGE_INTEGER freespace;
    BOOL bResult;
    bResult = GetDiskFreeSpaceExA(location, NULL, NULL, &freespace);
    if(!bResult)
    {
        CString szFormat;
        szFormat.Format(TEXT("Failed to get the free space on drive %s.\n"), location);
        CString szMsg = GetErrorMessage(szFormat);
        ::MessageBox(NULL, (LPCTSTR)szMsg, TEXT("File Error"), MB_OK);
        return TRUE;
    }
    return (spaceneeded <= freespace.QuadPart);
}

// given a drive letter (ending in a slash), return the label for that drive
// TODO make this more robust by adding input verification
CString GetDriveLabel(LPCTSTR szDev)
{
    CString retVal;
    TCHAR szVolumeNameBuffer[MAX_PATH + 1] = { 0 };
    BOOL bRet = ::GetVolumeInformation(szDev, szVolumeNameBuffer, MAX_PATH + 1, NULL,
                                        NULL, NULL, NULL, MAX_PATH + 1);

    // if malloc fails, nameBuf will be NULL.
    // if GetVolumeInfo fails, nameBuf will contain empty string
    // if all succeeds, nameBuf will contain label
    retVal = szVolumeNameBuffer;
    return retVal;
}

BOOL GetDisksProperty(HANDLE hDevice, PSTORAGE_DEVICE_DESCRIPTOR pDevDesc,
                      DEVICE_NUMBER *devInfo)
{
    STORAGE_PROPERTY_QUERY Query; // input param for query
    DWORD dwOutBytes; // IOCTL output length
    BOOL bResult; // IOCTL return val
    BOOL retVal = TRUE;
    DWORD cbBytesReturned;

    // specify the query type
    Query.PropertyId = StorageDeviceProperty;
    Query.QueryType = PropertyStandardQuery;

    // Query using IOCTL_STORAGE_QUERY_PROPERTY
    bResult = ::DeviceIoControl(hDevice, IOCTL_STORAGE_QUERY_PROPERTY,
                &Query, sizeof(STORAGE_PROPERTY_QUERY), pDevDesc,
                pDevDesc->Size, &dwOutBytes, (LPOVERLAPPED)NULL);
    if(bResult)
    {
        bResult = ::DeviceIoControl(hDevice, IOCTL_STORAGE_GET_DEVICE_NUMBER,
                    NULL, 0, devInfo, sizeof(DEVICE_NUMBER), &dwOutBytes,
                    (LPOVERLAPPED)NULL);
        if(!bResult)
        {
            retVal = FALSE;

            CString szMsg = GetErrorMessage(TEXT(
                "An error occurred while getting the device number.\n"
                "This usually means something is currently accessing the device;"
                "please close all applications and try again."));
            ::MessageBox(NULL, (LPCTSTR)szMsg, TEXT("File Error"), MB_OK);
        }
    }
    else
    {
        bResult = DeviceIoControl(hDevice, IOCTL_STORAGE_CHECK_VERIFY2, NULL, 0, NULL, 0, &cbBytesReturned,
                            (LPOVERLAPPED) NULL);
        if(bResult && GetLastError() == ERROR_INVALID_FUNCTION)
        {
            retVal = FALSE;
        }
        else
        {
            CString szMsg = GetErrorMessage(TEXT(
                "An error occurred while querying the properties.\n"
                "This usually means something is currently accessing the device;"
                "please close all applications and try again."));
            ::MessageBox(NULL, (LPCTSTR)szMsg, TEXT("File Error"), MB_OK);
        }
        retVal = FALSE;
    }

    return(retVal);
}

// some routines fail if there's no trailing slash in a name,
// 		others fail if there is.  So this routine takes a name (trailing
// 		slash or no), and creates 2 versions - one with the slash, and one w/o
//
// 		CALLER MUST FREE THE 2 RETURNED STRINGS
BOOL Slashify(char *str, char **slash, char **noSlash)
{
    BOOL retVal = FALSE;
    int strLen = strlen(str);
    if( strLen > 0 )
    {
        if( *(str + strLen - 1) == '\\' )
        {
            // trailing slash exists
            if(( (*slash = (char *)calloc( (strLen + 1), sizeof(char))) != NULL) &&
                    ( (*noSlash = (char *)calloc(strLen, sizeof(char))) != NULL))
            {
                strncpy(*slash, str, strLen);
                strncpy(*noSlash, *slash, (strLen - 1));
                retVal = TRUE;
            }
        }
        else
        {
            // no trailing slash exists
            if( ((*slash = (char *)calloc( (strLen + 2), sizeof(char))) != NULL) &&
                 ((*noSlash = (char *)calloc( (strLen + 1), sizeof(char))) != NULL) )
            {
                strncpy(*noSlash, str, strLen);
                sprintf(*slash, "%s\\", *noSlash);
                retVal = TRUE;
            }
        }
    }
    return retVal;
}

BOOL GetMediaType(HANDLE hDevice)
{
    DISK_GEOMETRY diskGeo;
    DWORD cbBytesReturned;
    if(DeviceIoControl(hDevice, IOCTL_DISK_GET_DRIVE_GEOMETRY,NULL, 0, &diskGeo, sizeof(diskGeo), &cbBytesReturned, NULL))
    {
        if((diskGeo.MediaType == FixedMedia) || (diskGeo.MediaType == RemovableMedia))
        {
            return TRUE; // Not a floppy
        }
    }
    return FALSE;
}


BOOL CheckDriveType(char *name, ULONG *pid)
{
    HANDLE hDevice;
    PSTORAGE_DEVICE_DESCRIPTOR pDevDesc;
    DEVICE_NUMBER deviceInfo;
    BOOL retVal = FALSE;
    char *nameWithSlash;
    char *nameNoSlash;
    int driveType;
    DWORD cbBytesReturned;

    // some calls require no tailing slash, some require a trailing slash...
    if(!(Slashify(name, &nameWithSlash, &nameNoSlash)) )
        return retVal;

    driveType = GetDriveTypeA(nameWithSlash);
    switch(driveType)
    {
    case DRIVE_REMOVABLE: // The media can be removed from the drive.
    case DRIVE_FIXED:     // The media cannot be removed from the drive. Some USB drives report as this.
        hDevice = CreateFileA(nameNoSlash, FILE_READ_ATTRIBUTES, FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, OPEN_EXISTING, 0, NULL);
        if(hDevice == INVALID_HANDLE_VALUE)
        {
            CString szFormat;
            szFormat.Format(TEXT("An error occurred when attempting to get a handle on %s.\n"), nameWithSlash);
            CString szMsg = GetErrorMessage(szFormat);
            ::MessageBox(NULL, (LPCTSTR)szMsg, TEXT("Volume Error"), MB_OK);

        }
        else
        {
            int arrSz = sizeof(STORAGE_DEVICE_DESCRIPTOR) + 512 - 1;
            pDevDesc = (PSTORAGE_DEVICE_DESCRIPTOR)new BYTE[arrSz];
            pDevDesc->Size = arrSz;

            // get the device number if the drive is
            // removable or (fixed AND on the usb bus, SD, or MMC (undefined in XP/mingw))
            if(GetMediaType(hDevice) && GetDisksProperty(hDevice, pDevDesc, &deviceInfo) &&
                    ( ((driveType == DRIVE_REMOVABLE) && (pDevDesc->BusType != BusTypeSata))
                      || ( (driveType == DRIVE_FIXED) && ((pDevDesc->BusType == BusTypeUsb)
                      || (pDevDesc->BusType == BusTypeSd ) || (pDevDesc->BusType == BusTypeMmc )) ) ) )
            {
                // ensure that the drive is actually accessible
                // multi-card hubs were reporting "removable" even when empty
                if(DeviceIoControl(hDevice, IOCTL_STORAGE_CHECK_VERIFY2, NULL, 0, NULL, 0, &cbBytesReturned, (LPOVERLAPPED) NULL))
                {
                    *pid = deviceInfo.DeviceNumber;
                    retVal = TRUE;
                }
                else
                // IOCTL_STORAGE_CHECK_VERIFY2 fails on some devices under XP/Vista, try the other (slower) method, just in case.
                {
                    CloseHandle(hDevice);
                    hDevice = CreateFileA(nameNoSlash, FILE_READ_DATA, FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, OPEN_EXISTING, 0, NULL);
                    if(DeviceIoControl(hDevice, IOCTL_STORAGE_CHECK_VERIFY, NULL, 0, NULL, 0, &cbBytesReturned, (LPOVERLAPPED) NULL))
                    {
                        *pid = deviceInfo.DeviceNumber;
                        retVal = TRUE;
                    }
                }
            }

            delete[] pDevDesc;
            CloseHandle(hDevice);
        }

        break;
    default:
        retVal = FALSE;
    }

    // free the strings allocated by slashify
    free(nameWithSlash);
    free(nameNoSlash);

    return retVal;
}
